# 在docker中拉取redis-stack ，然后再端口6379启动
# docker pull redis/redis-stack
# docker run -d -p 6379:6379 redis/redis-stack

# 安装相关库
# pip install langchain_redis redis 

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_community.llms import Tongyi
from langchain_redis import RedisChatMessageHistory
from dotenv import load_dotenv
import os

# 加载环境变量
load_dotenv()

# Redis连接配置
REDIS_URL = os.getenv("REDIS_URL", "redis://localhost:6379")
# print(f"连接到Redis: {REDIS_URL}")

# 初始化通义千问模型
api_key = os.getenv("DASHSCOPE_API_KEY")
llm = Tongyi(
    model="qwen-turbo",
    api_key=api_key,
    temperature=0.5,
    max_tokens=1024,
    max_retries=2
)

# 创建提示模板
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个友好的AI助手，能够记住之前的对话内容。"),
    MessagesPlaceholder(variable_name="history"),
    ("human", "{input}")
])

# 创建链
chain = prompt | llm

# Redis历史记录获取函数
def get_session_history(session_id: str) -> RedisChatMessageHistory:
    return RedisChatMessageHistory(
        session_id=session_id,
        redis_url=REDIS_URL,
        key_prefix="chat_history:"
    )

# 创建带历史记录的可运行对象
with_message_history = RunnableWithMessageHistory(
    chain,
    get_session_history,
    input_messages_key="input",
    history_messages_key="history",
)

# 测试对话
session_id = "user_redis_123"

print("\n=== Redis永久记忆对话测试 ===")

# 第一次对话
res1 = with_message_history.invoke(
    {"input": "你好！我是秀爷，我是一个抖音的视频创作者"},
    config={"configurable": {"session_id": session_id}}
)
print("\n第一次对话结果:")
print(res1)

# 第二次对话
res2 = with_message_history.invoke(
    {"input": "你还记得我是谁吗？我的职业是什么？"},
    config={"configurable": {"session_id": session_id}}
)
print("\n第二次对话结果:")
print(res2)

# 第三次对话（测试更复杂的记忆）
res3 = with_message_history.invoke(
    {"input": "我想要一些关于视频创作的建议"},
    config={"configurable": {"session_id": session_id}}
)
print("\n第三次对话结果:")
print(res3)

print("\n=== 对话历史已保存到Redis，重启程序后仍可恢复 ===")

